from utils import print_timing

@print_timing
def truncate_genofile(tmpdir, dataset, dbname):
    import os
    from utils import get_conf, pg_url
    conf = get_conf()
    dbuser = conf["dbuser"]
    password = conf["password"]
    portnum = conf["portnum"]
    dbstring = "%s://%s:%s@localhost:%s/%s"%(pg_url(), dbuser, password, portnum, dbname)
    from dbutils import get_db_type
    dbtype = get_db_type(dbstring)
    dirpath = os.path.dirname(os.path.abspath(__file__))
    datasets = conf[dbtype].sections
    if dataset not in datasets:
        sys.exit("error, dataset %s not listed in config"%dataset)

    annotfile = conf[dbtype]["annotfile"]
    phenofiles = conf[dbtype][dataset]["phenofiles"]
    genofile = conf[dbtype][dataset]["genofile"]
    test_genofile = os.path.basename(conf[dbtype][dataset]["genofile"]) + '.trunc'
    # since the schema is not actually used, just set schema to dataset
    schema = dataset
    if dbtype == 'affy6':
        from dbschema import make_tables_affy6
        table_dict = make_tables_affy6(schema)
        idlist = conf[dbtype]["fidlist"]
        from geno import Affy6
        g = Affy6(genofile, schema, dbstring, tmpdir, table_dict, 1)
    elif dbtype == 'illumina':
        from dbschema import make_tables_illumina
        table_dict = make_tables_illumina(schema)
        idlist = conf[dbtype]["rsidlist"]
        from geno import Illumina
        g = Illumina(genofile, schema, dbstring, tmpdir, table_dict, 1)
    try:
        g.truncate_genofile(test_genofile, idlist)
    except:
        # In case of error, rollback tmpdir creation
        import shutil
        if os.path.exists(tmpdir):
            shutil.rmtree(tmpdir)
        raise

if __name__=='__main__':
    import sys, os
    from optparse import OptionParser
    usr = os.popen('whoami').read().strip()
    usage = "'%prog tmpdir dataset dbname'. Using the value dbtype_" + usr + " for db is recommended, where dbtype can be either 'affy6' or 'illumina'."
    parser = OptionParser(usage=usage)
    (options, args) = parser.parse_args()
    if len(args) != 3:
        print "incorrect number of arguments"
        parser.print_help()
        sys.exit(1)
    tmpdir = args[0]
    dataset = args[1]
    dbname = args[2]
    truncate_genofile(tmpdir, dataset, dbname)
